'''
对特定类别进行重采样
'''
import os
import json
import shutil
from tqdm import tqdm

if __name__ == "__main__":
    imgpath = "Z:\\电科未检测到驾驶员视频附件_拆帧\\image_label"  #使用\\路径分隔
    jsonpath = "Z:\\电科未检测到驾驶员视频附件_拆帧\\json" #使用\\路径分隔
    resample_category_list = ["face", "hand", "cigarette", "cellphone"]
    resample_times = 2
    copy_img = True
    for jsonname in tqdm(os.listdir(jsonpath)):
        if jsonname.endswith(".json"):
            has_category = False
            imgname = jsonname.split(".")[0] + ".jpg"
            jsonfile = os.path.join(jsonpath, jsonname)
            imgfile = os.path.join(imgpath, imgname)
            with open(jsonfile, "r", encoding="utf-8") as f:
                json_content = json.load(f)
                for shape in json_content["shapes"]:
                    if shape["label"] in resample_category_list:
                        has_category = True
                        break
                if has_category:
                    for i in range(resample_times):
                        jsonfile_resample = jsonfile[:-5] + "_resample" + str(i+1) + ".json"
                        imgfile_resample = imgfile[:-4] + "_resample" + str(i+1) + ".jpg"
                        if not os.path.exists(jsonfile_resample):
                            with open(jsonfile_resample, "w", encoding="utf-8") as f1:
                                json_content["imagePath"] = imgfile_resample.split("\\")[-1]
                                json.dump(json_content, f1, indent=2, ensure_ascii=False)
                        if copy_img:
                            if not os.path.exists(imgfile_resample):
                                shutil.copy(imgfile, imgfile_resample)



